# -*- coding: utf-8 -*-
import pyodbc
import pymysql
import os
import pandas as pd
import argparse
import hashlib
from concurrent.futures import ThreadPoolExecutor, as_completed

# 多线程执行，数据查询
def parse_arguments():
    parser = argparse.ArgumentParser(description='Restore files based on database records')
    parser.add_argument('--database_type', required=True, choices=['mysql', 'sqlserver'],
                        help='Type of database (mysql or sqlserver)')
    parser.add_argument('--server', required=True, help='Database server address')
    parser.add_argument('--database', required=True, help='Database name')
    parser.add_argument('--username', required=True, help='Database username')
    parser.add_argument('--password', required=True, help='Database password')
    parser.add_argument('--base_directory', required=True, help='Base directory for file paths')
    parser.add_argument('--output_csv', required=True, help='Output CSV file path for missing files')
    parser.add_argument('--start_date', required=True, help='Start date for query (YYYY-MM-DD)')
    parser.add_argument('--end_date', required=True, help='End date for query (YYYY-MM-DD)')
    return parser.parse_args()


def query_mysql_database(server, database, username, password, start_date, end_date):
    """查询 MySQL 数据库并返回结果"""
    conn = pymysql.connect(host=server, user=username, passwd=password, db=database)
    cursor = conn.cursor()
    query = '''
    SELECT id, ABSOLUTEPATH, SYSFILENAME, OPERATIONCODE, OPERATIONID, PHYFILENAME, FILETYPE, FILESIZE
    FROM sys_accessories
    WHERE CREATETIME BETWEEN %s AND %s
    '''
    cursor.execute(query, (start_date, end_date))
    results = cursor.fetchall()
    conn.close()
    return results


def query_sqlserver_database(server, database, username, password, start_date, end_date):
    """查询 SQL Server 数据库并返回结果"""
    conn_str = (
        "DRIVER={{ODBC Driver 17 for SQL Server}};"
        "SERVER={};"
        "DATABASE={};"
        "UID={};"
        "PWD={}".format(server, database, username, password)
    )
    conn = pyodbc.connect(conn_str)
    cursor = conn.cursor()
    query = '''
    SELECT id, ABSOLUTEPATH, SYSFILENAME, OPERATIONCODE, OPERATIONID, PHYFILENAME, FILETYPE, FILESIZE
    FROM sys_accessories
    WHERE CREATETIME BETWEEN ? AND ?
    '''
    cursor.execute(query, (start_date, end_date))
    results = cursor.fetchall()
    conn.close()
    return results


def calculate_md5(file_path):
    """计算文件的 MD5 哈希值"""
    hash_md5 = hashlib.md5()
    with open(file_path, "rb") as f:
        for chunk in iter(lambda: f.read(4096), b""):
            hash_md5.update(chunk)
    return hash_md5.hexdigest()


def find_files_by_type_and_size(base_directory, file_type, file_size):
    """在指定目录中查找符合文件类型和文件大小的文件"""
    matched_files = []
    for root, _, files in os.walk(base_directory):
        for file in files:
            if file.endswith(file_type):
                file_path = os.path.join(root, file)
                if os.path.getsize(file_path) == file_size:
                    matched_files.append(file_path)
    return matched_files


def save_files_info_to_csv(results, base_directory, output_csv):
    """将恢复的文件信息保存到 CSV 文件"""
    data = []

    def process_result(result):
        file_id, abs_path, sys_filename, op_code, op_id, phy_filename, file_type, file_size = result
        matched_files = find_files_by_type_and_size(base_directory, file_type, file_size)
        for file_path in matched_files:
            file_md5 = calculate_md5(file_path)
            data.append({
                '操作代码': op_code,
                '业务ID': op_id,
                '文件ID': file_id,
                '文件名称': phy_filename,
                '文件路径': file_path,
                '文件MD5': file_md5
            })

    with ThreadPoolExecutor(max_workers=4) as executor:
        futures = [executor.submit(process_result, result) for result in results]
        for future in as_completed(futures):
            future.result()

    df = pd.DataFrame(data)

    # 打印 DataFrame 的列名和示例数据以调试
    print("DataFrame列名：", df.columns.tolist())
    print("DataFrame 示例数据：", df.head())

    # 保存 CSV 文件
    df.to_csv(output_csv, index=False, header=True)
    print("File information saved to {}".format(output_csv))


if __name__ == '__main__':
    args = parse_arguments()

    # 根据数据库类型查询数据库
    if args.database_type == 'mysql':
        results = query_mysql_database(args.server, args.database, args.username, args.password, args.start_date,
                                       args.end_date)
    elif args.database_type == 'sqlserver':
        results = query_sqlserver_database(args.server, args.database, args.username, args.password, args.start_date,
                                           args.end_date)

    # 将恢复的文件信息保存到 CSV 文件
    save_files_info_to_csv(results, args.base_directory, args.output_csv)

    print("Missing files saved to {}".format(args.output_csv))
